/*
 * Change Logs:
 * Date           Author       Notes
 * 2023-11-23     WangShun     the first version
 * 2024-09-19     WangShun     support rv32
 * 2024-10-11     zhangwk      support nuclei rv32(gcc and iar)  
 */

#include "../include/rvbacktrace.h"

/* Please check that the following symbols are defined in the linked scripts ！*/ 
/* If not, define the following symbols at the beginning and end of the text segment */
#define USE_GCC       1
#define USE_IAR       (!USE_GCC)

#if USE_GCC //nuclei gcc build
extern char *__etext;
extern char *__stext;
#define BT_CHK_PC_AVAIL(pc)   ((uintptr_t)(pc) < (uintptr_t)(&__etext) \
                              && (uintptr_t)(pc) > (uintptr_t)(&__stext))
#else //iar build, use icf, so we need set the var refer to the icf mem range 
char *__etext = (char *)0xB027FFFF;//sram1 cache end_addr
char *__stext = (char *)0xB0180000;//sram0 cache start_addr
#define BT_CHK_PC_AVAIL(pc)   ((uintptr_t)(pc) < (uintptr_t)(__etext) \
                              && (uintptr_t)(pc) > (uintptr_t)(__stext))

#define IAR_LIB_START_ADDR      0xB018F000                          
#define IAR_LIB_END_ADDR      0xB018FFFF   

#endif

extern unsigned int rvstack_frame[STACK_FRAME_LEN]; // stack frame
extern unsigned int rvstack_frame_len; // stack frame len

static int lvl;



#define BT_FUNC_LIMIT   0x2000
#define BT_LVL_LIMIT    64

#define BT_PC2ADDR(pc)        ((char *)(((uintptr_t)(pc))))

/* get framesize from c ins32 */
static int riscv_backtrace_framesize_get1(char** p_code_addr, unsigned int inst)
{
    unsigned int imm = 0;
    unsigned int lib_code_pc = 0;
    char *code_addr = *p_code_addr;

    /* addi sp, sp, -im */
    if ((inst & 0x800FFFFF) == 0x80010113) {
        imm = (inst >> 20) & 0x7FF;
        imm = (~imm & 0x7FF) + 1;
        return imm >> 2;//32bit rv32 sizet is 4, so move 2bits
    }

    /* jal      t0, __riscv_save_x*/
    if ((inst & 0x7F) == 0x6F) {
        imm = ((inst >> 21) & 0x3FF)<<1;//10:1
        imm |= (((inst >> 20) & 0x1) << 11);//bit11
        imm |=  (((inst >> 12) & 0xFF) << 12);//19:12

        //bit20 . The offffset is sign-extended 
        if(((inst >> 31) & 0x1) == 0)//bit20
        {
            lib_code_pc = ( unsigned int)code_addr + imm;
        }
        else
        {
            lib_code_pc = ( unsigned int)code_addr - (0x100000 - imm);// negative num need complement
        }
#if USE_IAR
        if((lib_code_pc < IAR_LIB_END_ADDR)&&(lib_code_pc > IAR_LIB_START_ADDR))
        {
            code_addr = (char *)(lib_code_pc);
            *p_code_addr = code_addr;//更新 p_code_addr，即 code_addr更新到 lib 里面
            return 0;//flag 0，即标志 for 循环重新找riscv_save 库里面的 c.swsp ra 指令
        }
#endif        
    }


    return -1;
}

/* get framesize from c ins */
static int riscv_backtrace_framesize_get(unsigned short inst)
{
    unsigned int imm = 0;
    /* addi sp, sp, -im */
    if ((inst & 0xFF83) == 0x1101) {
        imm = (inst >> 2) & 0x1F;
        imm = (~imm & 0x1F) + 1;
        return imm >> 2;//32bit rv32 sizet is 4, so move 2bits
    }

    /* c.addi16sp sp, nzuimm6 */
    if ((inst & 0xFF83) == 0x7101) {
        imm = (inst >> 3) & 0x3;
        imm = (imm << 1) | ((inst >> 5) & 0x1);
        imm = (imm << 1) | ((inst >> 2) & 0x1);
        imm = (imm << 1) | ((inst >> 6) & 0x1);
        imm = ((~imm & 0x1f) + 1) << 4;
        return imm >> 2;//32bit rv32 sizet is 4, so move 2bits
    }

    return -1;
}

static int riscv_backtrace_ra_offset_get1(unsigned int inst)
{
    unsigned int imm = 0;
    /* sw ra,552(sp) */
    if ((inst & 0x81FFF07F) == 0x112023) {
        imm = (inst >> 7) & 0x1F;
        imm |= ((inst >> 25) & 0x7F) << 5;
        return imm >> 2;//32bit rv32 sizet is 4, so move 2bits
    }

    return -1;
}

/* get ra position in the stack */
static int riscv_backtrace_ra_offset_get(unsigned short inst)
{
    unsigned int imm = 0;
    /* c.swsp  ra,0x3C(sp), uimm6<<2(sp) */
    if ((inst & 0xE07F) == 0xC006) {
        imm = (inst >> 7) & 0x3;//offset[7:6]
        imm = (imm << 6) | (((inst >> 9) & 0xf)<<2);
        return imm >> 2;//32bit rv32 sizet is 4, so move 2bits
    }

    return -1;
}
static char *k_int64tostr(int64_t num, char *str)
{
    char         index[] = "0123456789ABCDEF";
    unsigned int usnum   = (unsigned int)num;

    str[7] = index[usnum % 16];
    usnum /= 16;
    str[6] = index[usnum % 16];
    usnum /= 16;
    str[5] = index[usnum % 16];
    usnum /= 16;
    str[4] = index[usnum % 16];
    usnum /= 16;
    str[3] = index[usnum % 16];
    usnum /= 16;
    str[2] = index[usnum % 16];
    usnum /= 16;
    str[1] = index[usnum % 16];
    usnum /= 16;
    str[0] = index[usnum % 16];
    usnum /= 16;

    return str;
}

/* get the offset between the jump instruction and the return address */
static int backtraceFindLROffset(char *LR, int (*print_func)(const char *fmt, ...))
{
    int offset = 0;
    char *LR_indeed;
    unsigned int ins32;

    LR_indeed = BT_PC2ADDR(LR);

    /* Usually jump using the JAL instruction */
    ins32 = *(unsigned int *)(LR_indeed - 4);
    if ((ins32 & 0x3) == 0x3) {
        offset = 4;
    } else {
        offset = 2;
    }

    return offset;
}

static int riscv_backtraceFromStack(long **pSP, char **pPC,
                                    int (*print_func)(const char *fmt, ...))
{
    char *CodeAddr = NULL;
    long  *SP      = *pSP;
    char *PC       = *pPC;
    char *LR;
    int   i;
    int   framesize;
    int   offset = 0;
    unsigned int ins32;
    unsigned short ins16;

    unsigned int r_imm = 0;
    unsigned int r_ins32;
    unsigned short r_ins16;
    unsigned int lib_code_pc = 0;
    char *r_CodeAddr = NULL;

    /* 1. scan code, find lr pushed,     addi sp, sp, -im or c.addi16sp sp, nzuimm6  */
    for (i = 0; i < BT_FUNC_LIMIT;) {
        /* FIXME: not accurate from bottom to up. how to judge 2 or 4byte inst */
        //CodeAddr = (char *)(((long)PC & (~0x3)) - i);
        CodeAddr = (char *)(PC - i);
        ins32 = *(unsigned int *)(CodeAddr);
        if ((ins32 & 0x3) == 0x3) {
            ins16 = *(unsigned short *)(CodeAddr - 2);
            if ((ins16 & 0x3) != 0x3) {
                i += 4;
                framesize = riscv_backtrace_framesize_get1(&CodeAddr, ins32);
                if (framesize > 0) {
                    CodeAddr += 4;
                    break;
                }
                else if(framesize == 0)//flag 0，即标志 for 循环重新找riscv_save 库里面的 sw ra 指令
                {
                    PC =  (char *)(CodeAddr + 0x40);//pc need large than risc_save_x func range
                    i = 0;//i=0，且不 break,则 continue 即 for 循环重新再运行
                
                }
                continue;
            }
        }
        i += 2;
        ins16 = (ins32 >> 16) & 0xffff;
        framesize = riscv_backtrace_framesize_get(ins16);
        if (framesize >= 0) {
          
          r_ins16 = (*(unsigned short *)(CodeAddr-2))&0x3ff;//pre code
          
          if(r_ins16 == 0x2ef)//pre code is jal     
          {
            
            SP = SP + framesize;// recover addi sp, sp, -im
            
            r_ins32 = *(unsigned int *)(CodeAddr -2); //only need -2
            r_imm = ((r_ins32 >> 21) & 0x3FF)<<1;//10:1
            r_imm |= (((r_ins32 >> 20) & 0x1) << 11);//bit11
            r_imm |=  (((r_ins32 >> 12) & 0xFF) << 12);//19:12

            //bit20 . The offffset is sign-extended 
            if(((r_ins32 >> 31) & 0x1) == 0)//bit20
            {
                lib_code_pc = ( unsigned int)CodeAddr + r_imm;
            }
            else
            {
                lib_code_pc = ( unsigned int)CodeAddr -(0x100000 - r_imm);// negative num need complement
            }

#if USE_IAR
            if((lib_code_pc < IAR_LIB_END_ADDR)&&(lib_code_pc > IAR_LIB_START_ADDR))//  jal t0, __riscv_save_x
            {
                r_CodeAddr = (char *)(lib_code_pc);
                SP  += (0x40>>2);//recover sp change addi sp, sp -0x40
                 
                ins16 = *(unsigned short *)(r_CodeAddr);
                if(ins16 == 0x5301)//find c.li     t1, -0x20 in __riscv_save_4~7
                {
                    SP -= (0x20>>2);//recover sp change
                    PC =  (char *)(r_CodeAddr + 0x100);//pc need large than risc_save_x func range
                }
                else if(ins16 == 0x4301)//c.li     t1, 0
                {
                    SP -= 0;
                }
                else if(ins16 == 0x5341)//c.li     t1, -0x10
                {
                    SP -= (0x10>>2);//recover sp change
                    PC =  (char *)(r_CodeAddr + 0x100);//pc need large than risc_save_x func range
                }
                /* 2. scan code, find ins: c.swsp  ra,0x3C(sp) */
                for (i = 0; r_CodeAddr + i < PC;) {
                    ins32 = *(unsigned int *)(r_CodeAddr + i);
                    i += 2;
                    ins16 = ins32 & 0xffff;
                    offset = riscv_backtrace_ra_offset_get(ins16);
                    if (offset >= 0) {
                       /* 3. output */
                      LR     = (char *) * (SP -0x1);//(-0x40 + 0x3c)>>2=0x1   
                      break;
                    }
                }
                if (BT_CHK_PC_AVAIL(LR) == 0) {
                    if (print_func != NULL) {
                        print_func("End of stack backtracking\r\n");
                    }
                    return -1;
                }
                *pSP   = SP;
                rvstack_frame[lvl] = (unsigned int)(LR - 4);//jal t0 riscv_savex, offset=4

                print_func("[%d]Stack interval :[0x%016lx - 0x%016lx]  ra 0x%016lx pc 0x%016lx\n", lvl, SP, SP + framesize, LR, LR - 4);
                *pPC   = LR - 4;

                return 4 == 0 ? 1 : 0;
                break;
            }
#endif
          }
          else
          {
            CodeAddr += 2;
            ins16 = *(unsigned short *)(CodeAddr+2);
            if(ins16 == 0x5301)//find c.li     t1, -0x20 in __riscv_save_4~7
            {
                SP -= (0x20>>2);//recover sp change
            }
            else if(ins16 == 0x4301)//c.li     t1, 0
            {
                SP -= 0;
            }
            else if(ins16 == 0x5341)//c.li     t1, -0x10
            {
                 SP -= (0x10>>2);//recover sp change
            }
            break;
          }
        }
    }

    if (i == BT_FUNC_LIMIT) {
        /* error branch */
        if (print_func != NULL) {
            print_func("Backtrace fail!\r\n");
        }
        return -1;
    }

    /* 2. scan code, find ins: sw       ra, 0x48(sp) or c.swsp  ra,0x3C(sp) */
    for (i = 0; CodeAddr + i < PC;) {
        ins32 = *(unsigned int *)(CodeAddr + i);
        if ((ins32 & 0x3) == 0x3) {
            i += 4;
            offset = riscv_backtrace_ra_offset_get1(ins32);
            if (offset >= 0) {
                break;
            }
        } else {
            i += 2;
            ins16 = ins32 & 0xffff;
            offset = riscv_backtrace_ra_offset_get(ins16);
            if (offset >= 0) {
                break;
            }
        }
    }

    /* 3. output */
    LR     = (char *) * (SP + offset);

    if (BT_CHK_PC_AVAIL(LR) == 0) {
        if (print_func != NULL) {
            print_func("End of stack backtracking\r\n");
        }
        return -1;
    }
    *pSP   = SP + framesize;
    offset = backtraceFindLROffset(LR, print_func);

    rvstack_frame[lvl] = (unsigned int)(LR - offset);

    print_func("[%d]Stack interval :[0x%016lx - 0x%016lx]  ra 0x%016lx pc 0x%016lx\n", lvl, SP, SP + framesize, LR, LR - offset);
    *pPC   = LR - offset;

    return offset == 0 ? 1 : 0;
}

static int backtraceFromStack(long **pSP, char **pPC,
                              int (*print_func)(const char *fmt, ...))
{
    if (BT_CHK_PC_AVAIL(*pPC) == 0) {
        return -1;
    }

    return riscv_backtraceFromStack(pSP, pPC, print_func);
}

/* get the return address of the current function */
__attribute__((always_inline)) static inline void *backtrace_get_sp(void)
{
    void *sp;
    __asm__ volatile("mv %0, sp\n" : "=r"(sp));
    return sp;
}

/* get the return address of the current function */
__attribute__((always_inline)) static inline void *backtrace_get_pc(void)
{
    void *pc;
    __asm__ volatile("auipc %0, 0\n" : "=r"(pc));
    return pc;
}

/* printf call stack
   return levels of call stack */
int rvbacktrace_fomit(int (*print_func)(const char *fmt, ...))
{
    char *PC;
    long  *SP;
    int   ret;

    if (print_func == NULL) {
        print_func = printf;
    }

    SP = backtrace_get_sp();
    PC = backtrace_get_pc();

    print_func("\r\n---- Port by zhangwk, RV_Backtrace Call Frame Start: ----\r\n");
    print_func("###Please consider the value of ra as accurate and the value of sp as only for reference###\n");
    for (lvl = 0; lvl < BT_LVL_LIMIT; lvl++) {
        ret = backtraceFromStack(&SP, &PC, print_func);
        if (ret != 0) {
            rvstack_frame_len = lvl;
            break;
        }
    }
    rvbacktrace_addr2line((uint32_t *)&rvstack_frame[0], print_func);
    print_func("---- RV_Backtrace Call Frame End:----\r\n");
    print_func("\r\n");
    return lvl;
}
